package com.cg.ai.controller;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.*;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;
import org.springframework.core.io.ByteArrayResource;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.client.HttpServerErrorException;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.multipart.MultipartFile;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.cg.ai.service.ChatService;

import java.io.IOException;
import java.util.*;

@RequiredArgsConstructor
@RestController
@RequestMapping("/dify")
@Slf4j
public class DifyWorkflowController {

    private final RestTemplate restTemplate;
    private final ObjectMapper objectMapper;
    private final ChatService chatService;

    @Value("${dify.base-url}")
    private String difyBaseUrl;

    @Value("${dify.api-key}")
    private String difyApiKey;

    /**
     * 上传文件到 Dify
     */
    @PostMapping(value = "/upload", consumes = MediaType.MULTIPART_FORM_DATA_VALUE)
    public String upload(@RequestParam("file") MultipartFile file,
                         @RequestParam(value = "user", required = false) String userParam,
                         @RequestHeader(value = "X-User-Id", required = false) String userHeader) throws IOException {
        String user = resolveUser(userParam, userHeader);
        log.info("[Dify] 上传文件请求: name={}, size={} bytes", file.getOriginalFilename(), file.getSize());
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.MULTIPART_FORM_DATA);
        // 使用 Dify API Key 鉴权
        headers.set("Authorization", "Bearer " + difyApiKey);

        MultiValueMap<String, Object> body = new LinkedMultiValueMap<>();
        body.add("file", new ByteArrayResource(file.getBytes()) {
            @Override
            public String getFilename() {
                // 指定原始文件名，便于服务端识别扩展名
                return file.getOriginalFilename();
            }
        });
        body.add("user", user);

        HttpEntity<MultiValueMap<String, Object>> requestEntity = new HttpEntity<>(body, headers);
        String uploadUrl = trimTrailingSlash(difyBaseUrl) + "/files/upload";
        log.info("[Dify] 调用上传接口: url={}", uploadUrl);
        ResponseEntity<String> response = restTemplate.exchange(uploadUrl, HttpMethod.POST, requestEntity, String.class);
        log.info("[Dify] 上传响应: status={}, length={}", response.getStatusCodeValue(),
                response.getBody() == null ? 0 : response.getBody().length());
        return response.getBody();
    }

    /**
     * 运行工作流/对话
     */
    @PostMapping(value = "/workflows/run", consumes = MediaType.APPLICATION_JSON_VALUE)
    public Map<String, Object> runWorkFlow(@RequestBody(required = false) Map<String, Object> workFlowRunDto,
                              @RequestHeader(value = "X-User-Id", required = false) String userHeader,
                              @RequestParam(value = "user", required = false) String userParam) {
        String user = resolveUser(userParam, userHeader);
        Map<String, Object> payload = new HashMap<>();
        if (workFlowRunDto != null) payload.putAll(workFlowRunDto);
        payload.putIfAbsent("user", user);
        // 统一整理 inputs、query、文件字段，补齐默认值
        normalizeAndCompletePayload(payload);

        // 若显式传入 workflow_id，则走 Workflows API（触发+轮询）以避免长耗时导致的 504
        String workflowId = null;
        Object wfIdObj = payload.get("workflow_id");
        if (wfIdObj != null) {
            String s = String.valueOf(wfIdObj).trim();
            if (!s.isEmpty()) workflowId = s;
        }
        if (workflowId != null) {
            try {
                return runWorkflowWithPolling(workflowId, payload, user);
            } catch (Exception ex) {
                log.error("[Dify] 工作流触发/轮询失败: {}", ex.getMessage());
                Map<String, Object> err = new HashMap<>();
                err.put("answer", "");
                err.put("files", Collections.emptyList());
                err.put("error", "workflow_failed: " + ex.getMessage());
                return err;
            }
        }

        Object query = payload.get("query");
        log.info("[Dify] 运行工作流请求: payloadKeys={}, user={}", payload.keySet(), user);
        HttpHeaders headers = new HttpHeaders();
        // Dify 鉴权
        headers.set("Authorization", "Bearer " + difyApiKey);
        headers.setContentType(MediaType.APPLICATION_JSON);
        // 对于 chat-messages，采用 streaming 模式规避 504，并声明可接受 SSE
        Object respMode = payload.get("response_mode");
        if (!(respMode instanceof String s) || s.isBlank() || "blocking".equalsIgnoreCase(s)) {
            payload.put("response_mode", "streaming");
        }
        headers.setAccept(Arrays.asList(MediaType.TEXT_EVENT_STREAM, MediaType.APPLICATION_JSON));

        HttpEntity<Map<String, Object>> requestEntity = new HttpEntity<>(payload, headers);
        String runUrl = trimTrailingSlash(difyBaseUrl) + "/chat-messages";
        log.info("[Dify] 调用运行接口(advanced-chat): url={}, queryLen={}",
                runUrl, query == null ? 0 : String.valueOf(query).length());
        try {
            Long sid = getLongFromInputs(payload, "sessionId");
            Long uid = getLongFromInputs(payload, "userId");
            if (sid != null && uid != null && query != null) {
                // 保存用户指令到本地会话
                chatService.appendUserMessage(sid, uid, String.valueOf(query));
            }
        } catch (Exception e) {
            log.warn("[Dify] 记录用户指令消息失败: {}", e.getMessage());
        }
        ResponseEntity<String> response;
        try {
            response = restTemplate.exchange(runUrl, HttpMethod.POST, requestEntity, String.class);
        } catch (HttpServerErrorException e) {
            // 若网关 504/5xx，再次强制以 streaming 重试，尽量打破网关超时
            if (e.getStatusCode() != null && (e.getStatusCode().value() == 504 || e.getStatusCode().is5xxServerError())) {
                log.warn("[Dify] chat-messages 调用失败，状态={}，将以 streaming 重试", e.getStatusCode());
                payload.put("response_mode", "streaming");
                HttpEntity<Map<String, Object>> retryEntity = new HttpEntity<>(payload, headers);
                response = restTemplate.exchange(runUrl, HttpMethod.POST, retryEntity, String.class);
            } else {
                throw e;
            }
        }
        log.info("[Dify] 运行响应: status={}, length={}", response.getStatusCodeValue(),
                response.getBody() == null ? 0 : response.getBody().length());

        Map<String, Object> result = new HashMap<>();
        String respBody = response.getBody() == null ? "" : response.getBody();
        MediaType ct = response.getHeaders().getContentType();
        boolean looksLikeSse = (ct != null && MediaType.TEXT_EVENT_STREAM.includes(ct)) || respBody.startsWith("data:") || respBody.contains("\n\ndata:");
        if (looksLikeSse) {
            try {
                Map<String, Object> sseRes = parseSseResult(respBody);
                persistAssistantAnswerIfAny(payload, (String) sseRes.get("answer"));
                return sseRes;
            } catch (Exception e) {
                log.warn("[Dify] 解析SSE响应失败: {}", e.getMessage());
            }
        }
        try {
            JsonNode root = objectMapper.readTree(respBody.isEmpty() ? "{}" : respBody);
            JsonNode data = root.path("data").isMissingNode() ? root : root.path("data");
            String runId = data.path("id").asText(null);
            if (runId == null) runId = root.path("workflow_run_id").asText(null);
            JsonNode outputs = data.path("outputs");
            String answer = null;
            if (!outputs.isMissingNode()) {
                JsonNode answerNode = outputs.path("answer");
                if (!answerNode.isMissingNode() && !answerNode.isNull()) {
                    answer = answerNode.asText();
                }
            }
            if (answer == null) {
                // 兼容部分响应直接在根节点提供 answer
                JsonNode ansNode = root.path("answer");
                if (!ansNode.isMissingNode() && !ansNode.isNull()) answer = ansNode.asText();
            }
            Long sid = getLongFromInputs(payload, "sessionId");
            Long uid = getLongFromInputs(payload, "userId");
            if (sid != null && uid != null && answer != null && !answer.isBlank()) {
                // 保存助手回复到本地会话
                chatService.appendAssistantMessage(sid, uid, answer);
            }

            result.put("workflow_run_id", runId);
            result.put("answer", answer);
            if (!outputs.isMissingNode()) {
                JsonNode files = outputs.path("files");
                if (!files.isMissingNode() && files.isArray()) {
                    result.put("files", objectMapper.convertValue(files, List.class));
                } else {
                    result.put("files", Collections.emptyList());
                }
            } else {
                result.put("files", Collections.emptyList());
            }
        } catch (Exception e) {
            log.warn("[Dify] 解析工作流响应失败: {}", e.getMessage());
            // 尝试作为 SSE 再解析一次（有些服务会在 5xx 后仍返回 SSE 文本）
            try {
                Map<String, Object> sseParsed = parseSseResult(respBody);
                if (sseParsed != null) {
                    persistAssistantAnswerIfAny(payload, (String) sseParsed.get("answer"));
                    return sseParsed;
                }
            } catch (Exception ignored) {}
            result.put("answer", "");
            result.put("files", Collections.emptyList());
        }
        return result;
    }

    /**
     * 使用 Dify Workflows API：先触发运行，再轮询任务结果，避免网关 504。
     */
    private Map<String, Object> runWorkflowWithPolling(String workflowId, Map<String, Object> payload, String user) throws IOException {
        // 触发运行
        String runUrl = trimTrailingSlash(difyBaseUrl) + "/workflows/" + workflowId + "/run";
        HttpHeaders headers = new HttpHeaders();
        headers.set("Authorization", "Bearer " + difyApiKey);
        headers.setContentType(MediaType.APPLICATION_JSON);

        Map<String, Object> req = new HashMap<>();
        Map<String, Object> inputs = toMap(payload.get("inputs"));
        req.put("inputs", inputs == null ? Collections.emptyMap() : inputs);
        req.put("user", user);

        HttpEntity<Map<String, Object>> requestEntity = new HttpEntity<>(req, headers);
        log.info("[Dify] 触发工作流: url={}, inputsKeys={}", runUrl, (inputs == null ? Collections.emptySet() : inputs.keySet()));
        ResponseEntity<String> triggerResp = restTemplate.exchange(runUrl, HttpMethod.POST, requestEntity, String.class);
        String triggerBody = triggerResp.getBody() == null ? "{}" : triggerResp.getBody();

        JsonNode triggerJson;
        try {
            triggerJson = objectMapper.readTree(triggerBody);
        } catch (Exception e) {
            throw new IOException("invalid_trigger_response: " + e.getMessage());
        }

        // 兼容两种返回：task_id（需要轮询）或直接 data/outputs（阻塞完成）
        String taskId = triggerJson.path("task_id").asText(null);
        if (taskId == null || taskId.isBlank()) {
            // 尝试直接解析结果
            Map<String, Object> direct = parseWorkflowResult(triggerJson);
            persistAssistantAnswerIfAny(payload, (String) direct.get("answer"));
            return direct;
        }

        // 轮询任务结果
        long start = System.currentTimeMillis();
        long maxWaitMs = 110_000; // 110s，规避边缘网关 100s 超时
        long intervalMs = 1000;   // 1s
        String taskUrl = trimTrailingSlash(difyBaseUrl) + "/workflows/tasks/" + taskId;
        log.info("[Dify] 开始轮询任务: taskId={}", taskId);
        while (System.currentTimeMillis() - start < maxWaitMs) {
            ResponseEntity<String> pollResp = restTemplate.exchange(taskUrl, HttpMethod.GET,
                    new HttpEntity<>(headers), String.class);
            String body = pollResp.getBody() == null ? "{}" : pollResp.getBody();
            JsonNode json;
            try {
                json = objectMapper.readTree(body);
            } catch (Exception e) {
                throw new IOException("invalid_poll_response: " + e.getMessage());
            }
            String status = json.path("status").asText("");
            if ("succeeded".equalsIgnoreCase(status)) {
                Map<String, Object> ok = parseWorkflowResult(json);
                persistAssistantAnswerIfAny(payload, (String) ok.get("answer"));
                return ok;
            }
            if ("failed".equalsIgnoreCase(status) || "stopped".equalsIgnoreCase(status)) {
                Map<String, Object> err = new HashMap<>();
                err.put("answer", "");
                err.put("files", Collections.emptyList());
                err.put("error", "workflow_" + status);
                return err;
            }
            try { Thread.sleep(intervalMs); } catch (InterruptedException ignored) { Thread.currentThread().interrupt(); }
        }
        Map<String, Object> timeout = new HashMap<>();
        timeout.put("answer", "");
        timeout.put("files", Collections.emptyList());
        timeout.put("error", "workflow_timeout");
        return timeout;
    }

    /**
     * 解析工作流类响应（兼容触发返回或任务查询返回）。
     */
    private Map<String, Object> parseWorkflowResult(JsonNode root) {
        Map<String, Object> result = new HashMap<>();
        JsonNode data = root.path("data").isMissingNode() ? root : root.path("data");
        String runId = data.path("id").asText(null);
        if (runId == null) runId = root.path("workflow_run_id").asText(null);
        result.put("workflow_run_id", runId);

        String answer = null;
        JsonNode outputs = data.path("outputs");
        if (!outputs.isMissingNode()) {
            JsonNode answerNode = outputs.path("answer");
            if (!answerNode.isMissingNode() && !answerNode.isNull()) answer = answerNode.asText();
            JsonNode files = outputs.path("files");
            if (!files.isMissingNode() && files.isArray()) {
                result.put("files", objectMapper.convertValue(files, List.class));
            } else {
                result.put("files", Collections.emptyList());
            }
        } else {
            // 兼容直接在根节点提供 answer
            JsonNode ansNode = root.path("answer");
            if (!ansNode.isMissingNode() && !ansNode.isNull()) answer = ansNode.asText();
            result.put("files", Collections.emptyList());
        }
        result.put("answer", answer);
        return result;
    }

    /**
     * 解析 Dify chat-messages 的 SSE 响应，将 data: {...} 分片累加为最终 answer，并提取 files。
     */
    private Map<String, Object> parseSseResult(String sseText) throws IOException {
        if (sseText == null) sseText = "";
        String[] lines = sseText.split("\n");
        StringBuilder answerBuilder = new StringBuilder();
        List<Object> filesAgg = new ArrayList<>();
        for (String raw : lines) {
            if (raw == null) continue;
            String line = raw.trim();
            if (!line.startsWith("data:")) continue;
            String payload = line.substring(5).trim();
            if (payload.isEmpty()) continue;
            if ("[DONE]".equalsIgnoreCase(payload)) break;
            JsonNode node;
            try {
                node = objectMapper.readTree(payload);
            } catch (Exception ex) {
                // 跳过非 JSON 的分片
                continue;
            }
            // 优先 data.outputs.answer，其次 data.answer，再次根 answer
            String piece = null;
            JsonNode data = node.path("data");
            if (!data.isMissingNode()) {
                JsonNode outputs = data.path("outputs");
                if (!outputs.isMissingNode()) {
                    JsonNode ans = outputs.path("answer");
                    if (!ans.isMissingNode() && !ans.isNull()) piece = ans.asText();
                    JsonNode files = outputs.path("files");
                    if (!files.isMissingNode() && files.isArray()) {
                        List list = objectMapper.convertValue(files, List.class);
                        if (list != null) filesAgg.addAll(list);
                    }
                }
                if (piece == null) {
                    JsonNode ans2 = data.path("answer");
                    if (!ans2.isMissingNode() && !ans2.isNull()) piece = ans2.asText();
                }
            }
            if (piece == null) {
                JsonNode ans3 = node.path("answer");
                if (!ans3.isMissingNode() && !ans3.isNull()) piece = ans3.asText();
            }
            if (piece != null && !piece.isEmpty()) {
                answerBuilder.append(piece);
            }
        }
        Map<String, Object> result = new HashMap<>();
        result.put("answer", answerBuilder.toString());
        result.put("files", filesAgg);
        return result;
    }

    /**
     * 若存在会话与用户标识，则把助手的最终回答写入本地消息表。
     */
    private void persistAssistantAnswerIfAny(Map<String, Object> payload, String answer) {
        if (answer == null || answer.isBlank()) return;
        try {
            Long sid = getLongFromInputs(payload, "sessionId");
            Long uid = getLongFromInputs(payload, "userId");
            if (sid != null && uid != null) {
                chatService.appendAssistantMessage(sid, uid, answer);
            }
        } catch (Exception e) {
            log.warn("[Dify] 保存助手回复失败: {}", e.getMessage());
        }
    }

    /**
     * 去除 URL 末尾斜杠，避免重复拼接路径
     */
    private String trimTrailingSlash(String url) {
        if (url == null || url.isEmpty()) {
            return "";
        }
        return url.endsWith("/") ? url.substring(0, url.length() - 1) : url;
    }

    /**
     * 解析 user 标识：优先使用参数 user，其次使用请求头 X-User-Id，最后生成临时用户
     */
    private String resolveUser(String userParam, String userHeader) {
        String user = (userParam != null && !userParam.isEmpty()) ? userParam : userHeader;
        if (user == null || user.isEmpty()) {
            user = "web-user-" + System.currentTimeMillis();
        }
        return user;
    }

    /**
     * 规范化并补全 payload
     */
    private void normalizeAndCompletePayload(Map<String, Object> payload) {
        Object inputsObj = payload.get("inputs");
        if (inputsObj instanceof Map<?, ?> mapInputs) {
            Map<String, Object> inputs = (Map<String, Object>) mapInputs;
            Object inputVar = inputs.get("input");
            Map<String, Object> normalized = normalizeFileCandidate(inputVar);
            if (normalized == null) {
                Object files = payload.get("files");
                if (files instanceof List<?> fl && !fl.isEmpty()) {
                    normalized = normalizeFileCandidate(fl.get(0));
                } else if (files instanceof Map<?, ?> fMap) {
                    normalized = normalizeFileCandidate(fMap);
                }
                if (normalized == null) {
                    Object fileId = payload.get("file_id");
                    normalized = normalizeFileCandidate(fileId);
                }
            }
            if (normalized != null) {
                inputs.put("input", normalized);
            } else if (inputVar instanceof Map<?, ?> m) {
                Map<String, Object> ensured = mapFileIdToUploadFileId((Map<String, Object>) m);
                ensureTransferMethod(ensured);
                inputs.put("input", ensured);
            }
        }

        Object query = payload.get("query");
        if (!(query instanceof String s) || s.trim().isEmpty()) {
            Map<String, Object> inputs = toMap(payload.get("inputs"));
            if (inputs != null) {
                Object iv = inputs.get("input");
                if (iv instanceof String si && !si.isBlank()) {
                    query = si;
                }
            }
            if (!(query instanceof String qs) || qs.trim().isEmpty()) {
                query = "请根据文件内容进行分析并给出要点。";
            }
            payload.put("query", query);
        }

        payload.putIfAbsent("response_mode", "blocking");
    }

    /**
     * 若对象为 Map 则强转返回，否则返回 null
     */
    private Map<String, Object> toMap(Object obj) {
        return (obj instanceof Map<?, ?> m) ? (Map<String, Object>) m : null;
    }

    /**
     * 将“文件候选”统一为标准文件对象
     */
    private Map<String, Object> normalizeFileCandidate(Object candidate) {
        if (candidate == null) return null;
        if (candidate instanceof String s && !s.isBlank()) {
            Map<String, Object> fileObj = new HashMap<>();
            fileObj.put("type", "document");
            fileObj.put("upload_file_id", s.trim());
            fileObj.put("transfer_method", "local_file");
            return fileObj;
        }
        if (candidate instanceof List<?> list && !list.isEmpty()) {
            return normalizeFileCandidate(list.get(0));
        }
        if (candidate instanceof Map<?, ?> m) {
            Map<String, Object> fileObj = mapFileIdToUploadFileId((Map<String, Object>) m);
            ensureTransferMethod(fileObj);
            return fileObj;
        }
        return null;
    }

    /**
     * 从 payload.inputs 中读取 Long 值
     */
    private Long getLongFromInputs(Map<String, Object> payload, String key) {
        Map<String, Object> inputs = toMap(payload.get("inputs"));
        if (inputs == null) return null;
        Object val = inputs.get(key);
        if (val instanceof Number n) return n.longValue();
        if (val instanceof String s) {
            try { return Long.parseLong(s.trim()); } catch (NumberFormatException ignored) {}
        }
        return null;
    }

    /**
     * 确保文件对象包含 transfer_method 与 type 的默认值
     */
    private void ensureTransferMethod(Map<String, Object> fileObj) {
        if (fileObj == null) return;
        Object tm = fileObj.get("transfer_method");
        if (tm == null || String.valueOf(tm).isBlank()) {
            fileObj.put("transfer_method", "local_file");
        }
        if (!fileObj.containsKey("type")) {
            fileObj.put("type", "document");
        }
    }

    /**
     * 兼容旧字段：若仅提供 id，则映射为 upload_file_id
     */
    private Map<String, Object> mapFileIdToUploadFileId(Map<String, Object> fileObj) {
        if (fileObj == null) return null;
        Object uploadId = fileObj.get("upload_file_id");
        Object legacyId = fileObj.get("id");
        if ((uploadId == null || String.valueOf(uploadId).isBlank()) && legacyId != null && !String.valueOf(legacyId).isBlank()) {
            fileObj.put("upload_file_id", String.valueOf(legacyId));
        }
        return fileObj;
    }
}